import os
import shutil

def classify_files(source_dir="."):
    # 定义文件类型与目标文件夹的映射关系
    file_types = {
        "Text": [".txt", ".md", ".csv"],
        "Documents": [".pdf", ".docx", ".xlsx", ".pptx"],
        "Images": [".jpg", ".jpeg", ".png", ".gif", ".bmp"],
        "Archives": [".zip", ".rar", ".7z", ".tar"],
        "Videos": [".mp4", ".avi", ".mov", ".mkv"],
        "Audio": [".mp3", ".wav", ".flac", ".aac"],
        "Code": [".py", ".js", ".html", ".css", ".java"]
    }

    # 遍历源目录中的所有文件
    for filename in os.listdir(source_dir):
        file_path = os.path.join(source_dir, filename)
        
        # 跳过目录
        if os.path.isdir(file_path):
            continue

        # 获取文件扩展名
        _, ext = os.path.splitext(filename)
        ext = ext.lower()  # 统一转为小写

        # 查找对应的分类
        found = False
        for category, extensions in file_types.items():
            if ext in extensions:
                target_dir = os.path.join(source_dir, category)
                
                # 创建目标目录（如果不存在）
                if not os.path.exists(target_dir):
                    os.makedirs(target_dir)
                
                # 移动文件
                shutil.move(file_path, os.path.join(target_dir, filename))
                found = True
                break

        # 未分类文件处理
        if not found:
            other_dir = os.path.join(source_dir, "Other")
            if not os.path.exists(other_dir):
                os.makedirs(other_dir)
            shutil.move(file_path, os.path.join(other_dir, filename))

def get_directory():
    """获取用户输入的目录路径"""
    while True:
        path = input("请输入要整理的目录路径（直接回车使用当前目录）: ").strip()
        if not path:
            return os.getcwd()
        
        if os.path.isdir(path):
            return os.path.abspath(path)
        print(f"错误：'{path}' 不是有效目录，请重新输入")

if __name__ == "__main__":
    target_dir = get_directory()
    print(f"开始整理目录：{target_dir}")
    
    try:
        classify_files(target_dir)
        print("文件整理完成！")
    except Exception as e:
        print(f"整理过程中出现错误：{str(e)}")